/*
 * FeatureTensor.cpp
 *
 *  Created on: Dec 15, 2017
 *      Author: zy
 */

#include "include/FeatureTensor.h"
#include <iostream>
#include <vector>
#include "tensorflow/core/graph/default_device.h"

using namespace tensorflow;
using namespace std;

//#define TENSORFLOW_MODEL_META "x64/release/mars-small128.ckpt-68577.meta"
//#define TENSORFLOW_MODEL "./data/mars-small128.pb"
//#define TENSORFLOW_MODEL_META "x64/Release/1mars-small128.ckpt-68577.meta"
//#define TENSORFLOW_MODEL "x64/Release/1mars-small128.ckpt-68577"

FeatureTensor* FeatureTensor::instance = NULL;

FeatureTensor* FeatureTensor::getInstance(string model) {
	if (instance == NULL) {
		instance = new FeatureTensor(model);
	}
	return instance;
}

FeatureTensor::FeatureTensor(string model) {
	//prepare model:
	//model = "./model/tracker/mars-small128.pb";
	bool status = init(model);
	if (status == false)
	{
		std::cout << "init failed" << std::endl;
		cout << "press any key exit" << endl;
		getchar();
		exit(1);
	}
	else {
		std::cout << "init succeed" << std::endl;
	}
}

FeatureTensor::~FeatureTensor() {
	session->Close();
	delete session;
	output_tensors.clear();
	outnames.clear();
}

bool FeatureTensor::init(string model) {
	tensorflow::SessionOptions sessOptions;
	//tensorflow::ConfigProto configProto;
	tensorflow::GPUOptions* gpuOptions = new tensorflow::GPUOptions;

	//tensorflow::ConfigProto* config = &sessOptions.config;


	gpuOptions->set_per_process_gpu_memory_fraction(0.65);
	gpuOptions->set_allow_growth(true);



	sessOptions.config.set_allocated_gpu_options(gpuOptions);


	Status status = NewSession(sessOptions, &session);

	if (session == nullptr) return false;

	const tensorflow::string pathToGraph = model;
	//tensorflow::MetaGraphDef graph_def;  // 原始
	tensorflow::GraphDef graph_def;  //不同类型的文件不同的定义

	status = ReadBinaryProto(tensorflow::Env::Default(), pathToGraph, &graph_def);
	if (status.ok() == false) {
		cout << status << endl;
		return false;
	}
	//status = session->Create(graph_def.graph_def());  // 原始
	status = session->Create(graph_def);
	if (status.ok() == false) {
		cout << status << endl;
		return false;
	}
	//std::vector<std::string> node_names;
	/*for (const auto &node : graph_def.graph_def().node()) {
		printf("node name:%s\n", node.name().c_str());
	}*/

	/*const tensorflow::string checkpointPath = TENSORFLOW_MODEL;
	Tensor checkpointTensor(DT_STRING, TensorShape());
	checkpointTensor.scalar<std::string>()() = checkpointPath;*/

	/*cout << graph_def.mutable_saver_def()->filename_tensor_name() << endl;
	cout << graph_def.mutable_saver_def()->restore_op_name() << endl;
	status = session->Run(
		{ { graph_def.mutable_saver_def()->filename_tensor_name(), checkpointTensor }, },
		{}, { graph_def.mutable_saver_def()->restore_op_name() }, nullptr);*/
	Tensor input_tensor(DT_UINT8, TensorShape({ 1, 128, 64, 3 }));
	input_layer = "images:0";
	outnames.push_back("Cast:0");
	std::vector<std::pair<tensorflow::string, Tensor>> feed_dict = {
			{input_layer, input_tensor},
	};

	status = session->Run(feed_dict, outnames, {}, &output_tensors);
	//cout << graph_def.saver_def().filename_tensor_name() << endl;
	//cout << graph_def.saver_def().restore_op_name() << endl;
	//status = session->Run(
	//		{ {graph_def.saver_def().filename_tensor_name(), checkpointTensor}, },
	//		{}, {graph_def.saver_def().restore_op_name()}, nullptr );  // 原始


	//string inputLayer = "input:0";
	//vector<string> outputLayer;  //注意，在yolo3中输出的三个节点名称为以下三个
	//outputLayer.push_back("output/boxes:0");
	//outputLayer.push_back("output/scores:0");
	//outputLayer.push_back("output/labels:0");
	//std::vector<std::pair<string, Tensor>> inputs = {
	//   { inputLayer, checkpointTensor }
	//};
	//std::vector<Tensor> outputs;
	//status = session->Run({ inputs }, { outputLayer }, {}, &outputs);
	if (status.ok() == false) {
		cout << status << endl;
		return false;
	}

	//input_layer = "placeholder:0";
	/*input_layer = "images:0";
	outnames.push_back("truediv:0");*/
	/*input_layer = "save/Const:0";
	outnames.push_back("Cast:0");*/
	feature_dim = 128;


	return true;
}

bool FeatureTensor::getRectsFeature(const cv::Mat& img, DETECTIONS& d) {
	std::vector<cv::Mat> mats;
	for (DETECTION_ROW& dbox : d) {
		cv::Rect rc = cv::Rect(int(dbox.tlwh(0)), int(dbox.tlwh(1)),
			int(dbox.tlwh(2)), int(dbox.tlwh(3)));
		rc.x -= (rc.height * 0.5 - rc.width) * 0.5;
		rc.width = rc.height * 0.5;
		rc.x = (rc.x >= 0 ? rc.x : 0);
		rc.y = (rc.y >= 0 ? rc.y : 0);
		rc.width = (rc.x + rc.width <= img.cols ? rc.width : (img.cols - rc.x));
		rc.height = (rc.y + rc.height <= img.rows ? rc.height : (img.rows - rc.y));

		cv::Mat mattmp = img(rc).clone();
		//cv::Mat mattmp = img.clone();
		cv::resize(mattmp, mattmp, cv::Size(64, 128));
		mats.push_back(mattmp);
	}
	int count = mats.size();

	/*Tensor input_tensor(DT_UINT8, TensorShape({count, 128, 64, 3}));
	tobuffer(mats, input_tensor.flat<uint8>().data());*/
	Tensor input_tensor(DT_UINT8, TensorShape({ count, 128, 64, 3 }));
	tobuffer(mats, input_tensor.flat<uint8>().data());
	std::vector<std::pair<tensorflow::string, Tensor>> feed_dict = {
			{input_layer, input_tensor},
	};

	/*cout << feed_dict[0] << endl;
	cout << outnames << endl;*/
	Status status = session->Run(feed_dict, outnames, {}, &output_tensors);
	if (status.ok() == false) {
		cout << status << endl;
		return false;
	}

	float* tensor_buffer = output_tensors[0].flat<float>().data();
	int i = 0;
	for (DETECTION_ROW& dbox : d) {
		for (int j = 0; j < feature_dim; j++)
			dbox.feature[j] = tensor_buffer[i * feature_dim + j];
		i++;
	}
	return true;
}

void FeatureTensor::tobuffer(const std::vector<cv::Mat>& imgs, uint8* buf) {
	int pos = 0;
	for (const cv::Mat& img : imgs) {
		int Lenth = img.rows * img.cols * 3;
		int nr = img.rows;
		int nc = img.cols;
		if (img.isContinuous()) {
			nr = 1;
			nc = Lenth;
		}
		for (int i = 0; i < nr; i++) {
			const uchar* inData = img.ptr<uchar>(i);
			for (int j = 0; j < nc; j++) {
				buf[pos] = *inData++;
				pos++;
			}
		}//end for
	}//end imgs;
}
void FeatureTensor::test()
{
	return;
}
